from fastapi import APIRouter, HTTPException, Depends
from schemas import StockResponse, UserOut
from typing import Optional
import pandas as pd
from dependencies import get_current_user
from utils import load_parquet

def get_stock_router(app):
    """Create and configure stock router with access to FastAPI app"""
    router = APIRouter(prefix="/api/stocks", tags=["Stocks"])

    @router.get("/{stock_code}", response_model=StockResponse)
    async def get_stock_data(
        stock_code: str,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
        current_user: UserOut = Depends(get_current_user)
    ):
        """Get stock data for a specific stock code"""
        try:
            df = app.state.df
            if df is None:
                raise HTTPException(status_code=500, detail="后端数据未加载")

            # Filter by stock code
            filtered_df = df[df['code'] == stock_code].copy()

            if filtered_df.empty:
                raise HTTPException(status_code=404, detail="股票代码不存在")

            # Date filtering (using date column, format YYYY-MM-DD)
            if 'date' in filtered_df.columns:
                if start_date:
                    try:
                        pd.to_datetime(start_date)  # Validate format
                        filtered_df = filtered_df[filtered_df['date'] >= start_date]
                    except ValueError:
                        raise HTTPException(status_code=400, detail="start_date 格式错误，应为YYYY-MM-DD")
                if end_date:
                    try:
                        pd.to_datetime(end_date)  # Validate format
                        filtered_df = filtered_df[filtered_df['date'] <= end_date]
                    except ValueError:
                        raise HTTPException(status_code=400, detail="end_date 格式错误，应为YYYY-MM-DD")
            else:
                raise HTTPException(status_code=500, detail="Parquet 文件中缺少 date 列")

            # Final formatting
            filtered_df = filtered_df.sort_values('date')
            filtered_df.replace({pd.NA: None, float('nan'): None}, inplace=True)

            return {
                "code": stock_code,
                "data": filtered_df.to_dict('records')
            }

        except HTTPException as http_exc:
            raise http_exc
        except Exception as e:
            print(f"获取股票数据时发生错误: {e}")
            raise HTTPException(status_code=500, detail=str(e))

    @router.get("/codes", response_model=dict)
    async def get_stock_codes(current_user: UserOut = Depends(get_current_user)):
        """Get list of all stock codes"""
        try:
            df = app.state.df
            if df is None:
                raise HTTPException(status_code=500, detail="后端数据未加载")
            codes = df['code'].unique().tolist()
            return {"codes": sorted(codes)}
        except HTTPException as http_exc:
            raise http_exc
        except Exception as e:
            print(f"获取股票代码列表时发生错误: {e}")
            raise HTTPException(status_code=500, detail=str(e))

    return router